/*
  Author: Chen Haodong (shinkunyachen@outlook.com)
  -----
  Copyright (c)2020 - present Chen Haodong. All Rights Reserved.
 */

#include "DetectMe/DetectMe.h"
#include "DetectMe/OperationDetection.h"
#include "DetectMe/Expression.h"
extern Report report;
// helper
namespace ispointer {
static bool is_pointer(const clang::Expr *expr) {
  switch (expr->getStmtClass()) {
  case clang::Stmt::DeclRefExprClass:
    return llvm::cast<clang::DeclRefExpr>(expr)
        ->getDecl()
        ->getType()
        .getTypePtr()
        ->isPointerType();
  case clang::Stmt::MemberExprClass:
    return ispointer::is_pointer(
        llvm::cast<clang::MemberExpr>(expr)->getBase());
  default:
    return false;
  }
}
} // namespace ispointer
static const clang::Stmt *is_in_loop(const clang::Stmt *stmt,
                                     clang::ASTContext &ctx) {
  if (!stmt)
    return nullptr;
  auto it = ctx.getParents(*stmt).begin();
  if (it == ctx.getParents(*stmt).end())
    return nullptr;
  const clang::Stmt *parent = it->get<clang::Stmt>();
  if (parent) {
    switch (parent->getStmtClass()) {
    case clang::Stmt::WhileStmtClass:
    case clang::Stmt::DoStmtClass:
    case clang::Stmt::ForStmtClass:
      return parent;
    default:
      break;
    }
    const clang::Stmt *result = is_in_loop(parent, ctx);
    if (result)
      return result;
  }
  return nullptr;
}
static void dispatch_set_value_holder(UsePoint *up, const clang::Expr *expr) {
  switch (expr->getStmtClass()) {
  case clang::Stmt::DeclRefExprClass:
    up->set_value_holder<ClangDeclRefExprValueHolder>(
        llvm::cast<clang::DeclRefExpr>(expr));
    break;
  case clang::Stmt::MemberExprClass:
    up->set_value_holder<ClangMemberExprValueHolder>(
        llvm::cast<clang::MemberExpr>(expr));
    break;
  default:
    break;
  }
}
void DetectMe::check() {
  auto ast_functions = resource->getFunctions(true);
  LangOptions LangOpts;
  LangOpts.CPlusPlus = true;
  for (auto f : ast_functions) {
    clang::FunctionDecl *func = manager->getFunctionDecl(f);
    //std::string name = func->getQualifiedNameAsString();
    //llvm::errs() << "Analyzing " << name << '\n';
    string filename = "";
    const clang::SourceManager &SM = func->getASTContext().getSourceManager();
    SourceLocation SL = func->getBeginLoc();
    filename = SL.printToString(SM);
    if (filename.find(".h") != string::npos || filename.find("include") != string::npos)
            continue;
    std::unique_ptr<CFG> &cfg = manager->getCFG(f);
    if (cfg) {
      WalkContext ctx(cfg.get());
      ctx.function = f;
      for (auto blk_it = cfg.get()->begin(), blk_end = cfg.get()->end();
           blk_it != blk_end; ++blk_it)
        walk(*blk_it, ctx);
    }
  }
}

void DetectMe::walk(const clang::CFGBlock *blk, WalkContext &ctx) {
  ctx.block = blk;
  for (auto it = blk->rbegin(), end = blk->rend(); it != end; ++it) {
    if (it->getKind() == clang::CFGElement::Statement) {
      CFGStmt s = it->castAs<CFGStmt>();
      // s.getStmt()->dump();
      ctx.stmt_it = it;
      walk(s.getStmt(), ctx);
    }
  }
}

void DetectMe::walk(const clang::Stmt *stmt, WalkContext &ctx) {
  if (stmt->getStmtClass() >= clang::Stmt::firstExprConstant &&
      stmt->getStmtClass() <= clang::Stmt::lastExprConstant) {
    walk(llvm::cast<clang::Expr>(stmt), ctx);
  } else if (stmt->getStmtClass() == clang::Stmt::DeclStmtClass) {
    const clang::DeclStmt *decl_stmt = llvm::cast<clang::DeclStmt>(stmt);
    const clang::Decl *decl = decl_stmt->getSingleDecl();
    if (decl->getKind() == clang::Decl::Var) {
      const clang::VarDecl *var_decl = llvm::cast<clang::VarDecl>(decl);
      const clang::Expr *init_expr = var_decl->getInit();
      if (init_expr) {
        walk(init_expr, ctx);
      }
    }
  }
  // TODO more stmt
}
void DetectMe::walk(const clang::Expr *expr, WalkContext &ctx) {
  // expr->dump();
  // llvm::errs() << '\n';
  // detect
  switch (expr->getStmtClass()) {
  case clang::Stmt::CompoundAssignOperatorClass:
  case clang::Stmt::BinaryOperatorClass: {
    auto e = llvm::cast<clang::BinaryOperator>(expr);
    detect_pointer_arithmetic_oob(e, ctx);
    detect_divmod_zero(e, ctx);
    detect_arithmetic_overflow(e, ctx);
    detect_infinite_loop(e, ctx);
    // detect_infinite_loop(e,ctx);
    break;
  }
  case clang::Stmt::UnaryOperatorClass: {
    auto e = llvm::cast<clang::UnaryOperator>(expr);
    detect_arithmetic_overflow(e, ctx);
    detect_null_dereference(e, ctx);
    detect_use_after_free(e, ctx);
    detect_infinite_loop(e, ctx);
    break;
  }
  case clang::Stmt::ArraySubscriptExprClass: {
    auto e = llvm::cast<clang::ArraySubscriptExpr>(expr);
    detect_array_oob_in_loop(e, ctx);
    break;
  }
  case clang::Stmt::MemberExprClass: {
    auto e = llvm::cast<clang::MemberExpr>(expr);
    detect_cast_oob(e, ctx);
    break;
  }
  // TODO more case
  default:
    break;
  }
  // walk
  switch (expr->getStmtClass()) {
  case clang::Stmt::CompoundAssignOperatorClass:
  case clang::Stmt::BinaryOperatorClass: {
    auto e = llvm::cast<clang::BinaryOperator>(expr);
    walk(e->getLHS(), ctx);
    walk(e->getRHS(), ctx);
    break;
  }
  case clang::Stmt::CallExprClass: {
    auto e = llvm::cast<clang::CallExpr>(expr);
    for (int i = 0, sz = e->getNumArgs(); i != sz; ++i)
      walk(e->getArg(i), ctx);
    break;
  }
  case clang::Stmt::UnaryOperatorClass: {
    walk(llvm::cast<clang::UnaryOperator>(expr)->getSubExpr(), ctx);
    break;
  }
  // TODO more case
  default:
    if (expr->getStmtClass() >= clang::Stmt::firstCastExprConstant &&
        expr->getStmtClass() <= clang::Stmt::lastCastExprConstant) {
      walk(llvm::cast<clang::CastExpr>(expr)->getSubExpr(), ctx);
    }
    break;
  }
}
void DetectMe::reportJson(const clang::Expr *expr, WalkContext &ctx,
                          DefectType dt) {
  ASTFunction *AF = ctx.function;
  FunctionDecl *FD = manager->getFunctionDecl(AF);
  SourceManager &SM = FD->getASTContext().getSourceManager();
  SourceLocation SL = expr->getBeginLoc();
  unsigned int lineNum = SM.getExpansionLineNumber(SL);
  string filename = SL.printToString(SM);
  report.addToReport(dt, filename, lineNum);
}
void DetectMe::detect_pointer_arithmetic_oob(const clang::BinaryOperator *expr,
                                             WalkContext &ctx) {
  switch (expr->getOpcode()) {
  case clang::BinaryOperatorKind::BO_Add: {
    const clang::Expr *lhs = extract_simple_value_expr(expr->getLHS());
    if (lhs) {
      if (ispointer::is_pointer(lhs)) {
        OperationDetection od(manager, call_graph, ctx.function);
        od.set_operation<BranchOperation>();
        od.set_use_point(ctx.block, ctx.stmt_it);
        dispatch_set_value_holder(od.use_point(), lhs);
        if (!od.detect(OperationDetection::Forall)) {
          DefectType dt = DefectType::PointerOutOfBound;
          reportJson(lhs, ctx, dt);
          llvm::errs() << "Defect1 Pointer Arithmetic OOB\n";
          expr->dump();
          llvm::errs() << '\n';
        }
      }
    }
    break;
  }
  case clang::BinaryOperatorKind::BO_AddAssign: {
    const clang::Expr *lhs = extract_simple_value_expr(expr->getLHS());
    if (lhs) {
      if (ispointer::is_pointer(lhs)) {
        OperationDetection od(manager, call_graph, ctx.function);
        od.set_operation<BranchOperation>();
        od.set_use_point(ctx.block, ctx.stmt_it);
        dispatch_set_value_holder(od.use_point(), lhs);
        if (!od.detect(OperationDetection::Forall)) {
          DefectType dt = DefectType::PointerOutOfBound;
          reportJson(lhs, ctx, dt);
          llvm::errs() << "Defect1 Pointer Arithmetic OOB\n";
          expr->dump();
          llvm::errs() << '\n';
        }
      }
    }
    break;
  }
  default:
    break;
  }
}
void DetectMe::detect_array_oob_in_loop(const clang::ArraySubscriptExpr *expr,
                                        WalkContext &ctx) {
  if (is_in_loop(expr,
                 manager->getFunctionDecl(ctx.function)->getASTContext())) {
    std::vector<const clang::Expr *> decl_exprs;
    collect_expr(expr->getBase(), clang::Stmt::DeclRefExprClass, decl_exprs);
    if (decl_exprs.size() > 0) {
      const clang::DeclRefExpr *decl_expr =
          llvm::cast<clang::DeclRefExpr>(decl_exprs[0]);
      OperationDetection od(manager, call_graph, ctx.function);
      od.set_operation<BranchOperation>();
      od.set_use_point(ctx.block, ctx.stmt_it);
      dispatch_set_value_holder(od.use_point(), decl_expr);
      if (!od.detect(OperationDetection::Forall)) {
        DefectType dt = DefectType::LoopWithExternalInput;
        reportJson(decl_expr, ctx, dt);
        llvm::errs() << "Defect2 Array OOB In Loop\n";
        expr->dump();
        llvm::errs() << '\n';
      }
    }
  }
}
void DetectMe::detect_cast_oob(const clang::MemberExpr *expr,
                               WalkContext &ctx) {
  std::vector<const clang::Expr *> decl_exprs;
  collect_expr(expr, clang::Stmt::DeclRefExprClass, decl_exprs);
  if (decl_exprs.size() > 0) {
    const clang::DeclRefExpr *decl_expr =
        llvm::cast<clang::DeclRefExpr>(decl_exprs[0]);
    OperationDetection od(manager, call_graph, ctx.function);
    od.set_operation<ExplicitCastOperation>();
    od.set_use_point(ctx.block, ctx.stmt_it);
    dispatch_set_value_holder(od.use_point(), decl_expr);
    if (od.detect(OperationDetection::Exists)) {
      OperationDetection od2(manager, call_graph, ctx.function);
      od2.set_operation<BranchOperation>();
      od2.set_use_point(ctx.block, ctx.stmt_it);
      dispatch_set_value_holder(od2.use_point(), decl_expr);
      if (!od2.detect(OperationDetection::Forall)) {
        DefectType dt = DefectType::StructOutOfBound;
        reportJson(decl_expr, ctx, dt);
        llvm::errs() << "Defect3 Cast OOB\n";
        expr->dump();
        llvm::errs() << '\n';
      }
    }
  }
}
// void DetectMe::detect_dereference_oob_after_arithmetic_overflow()// unknown
void DetectMe::detect_infinite_loop(const clang::Expr *expr, WalkContext &ctx) {
  /* const clang::Expr *suspect = nullptr;
  const clang::Expr *need_check_expr = nullptr;
  switch (expr->getStmtClass()) {
  case clang::Stmt::CompoundAssignOperatorClass:
  case clang::Stmt::BinaryOperatorClass: {
    const clang::BinaryOperator *bexpr =
        llvm::cast<clang::BinaryOperator>(expr);
    suspect = extract_simple_value_expr(bexpr->getLHS());
    switch (bexpr->getOpcode()) {
    case clang::BinaryOperatorKind::BO_AddAssign:
    case clang::BinaryOperatorKind::BO_SubAssign:
    case clang::BinaryOperatorKind::BO_MulAssign:
      need_check_expr = expr;
      break;
    case clang::BinaryOperatorKind::BO_Assign:
      need_check_expr = bexpr->getRHS();
      break;
      default:
        break;
      }
    break;
  }
  case clang::Stmt::UnaryOperatorClass: {
    const clang::UnaryOperator *uexpr = llvm::cast<clang::UnaryOperator>(expr);
    suspect = extract_simple_value_expr(uexpr->getSubExpr());
    need_check_expr = expr;
    break;
  }
  default:
    break;
  }
  if (!suspect || !need_check_expr)
    return;
  const clang::Stmt *loop =
      is_in_loop(expr, manager->getFunctionDecl(ctx.function)->getASTContext());
  if (loop) {
    // get condition
    const clang::Expr *cond = nullptr;
    switch (loop->getStmtClass()) {
    case clang::Stmt::WhileStmtClass:
      cond = llvm::cast<clang::WhileStmt>(loop)->getCond();
      break;
    case clang::Stmt::DoStmtClass:
      cond = llvm::cast<clang::DoStmt>(loop)->getCond();
      break;
    case clang::Stmt::ForStmtClass:
      cond = llvm::cast<clang::ForStmt>(loop)->getCond();
      break;
    default:
      break;
    }
    if (cond) {
      std::vector<const clang::Expr *> exprs;
      collect_expr(cond, clang::Stmt::DeclRefExprClass, exprs);
      collect_expr(cond, clang::Stmt::MemberExprClass, exprs);
      for (auto e : exprs) {
        if (is_equal(e, suspect)) {
          if (detect_arithmetic_overflow(need_check_expr, ctx)) {
            DefectType dt = DefectType::DeadLoop;
            reportJson(suspect, ctx, dt);
            llvm::errs() << "Defect5 Infinite Loop\n";
            expr->dump();
            llvm::errs() << '\n';
          }
        }
      }
    }
  } */
}
void DetectMe::detect_divmod_zero(const clang::BinaryOperator *expr,
                                  WalkContext &ctx) {
  switch (expr->getOpcode()) {
  case clang::BinaryOperatorKind::BO_Div:
  case clang::BinaryOperatorKind::BO_Rem:
  case clang::BinaryOperatorKind::BO_DivAssign:
  case clang::BinaryOperatorKind::BO_RemAssign: {
    const clang::Expr *rhs = extract_simple_value_expr(expr->getRHS());
    if (rhs) {
      OperationDetection od(manager, call_graph, ctx.function);
      od.set_operation<BranchOperation>();
      od.set_use_point(ctx.block, ctx.stmt_it);
      dispatch_set_value_holder(od.use_point(), rhs);
      if (!od.detect(OperationDetection::Forall)) {
        DefectType dt = DefectType::DivOrDodByZero;
        reportJson(rhs, ctx, dt);
        llvm::errs() << "Defect6 Div/Mod Zero\n";
        expr->dump();
        llvm::errs() << '\n';
      }
    } else {
      const clang::Expr *rhs_expr = expr->getRHS();
      if (rhs_expr->getStmtClass() == clang::Stmt::IntegerLiteralClass) {
        const clang::IntegerLiteral *integer =
            llvm::cast<clang::IntegerLiteral>(rhs_expr);
        if (integer->getValue().isNullValue()) {
          DefectType dt = DefectType::DivOrDodByZero;
          reportJson(rhs, ctx, dt);
          llvm::errs() << "Defect6 Div/Mod Zero\n";
          expr->dump();
          llvm::errs() << '\n';
        }
      }
    }
    break;
  }
  default:
    break;
  }
}
// void detect_sql_inject(); // unknown
bool DetectMe::detect_arithmetic_overflow(const clang::Expr *stmt,
                                          WalkContext &ctx) {
  if (stmt->getStmtClass() == clang::Stmt::BinaryOperatorClass ||
      stmt->getStmtClass() == clang::Stmt::CompoundAssignOperatorClass) {
    const clang::BinaryOperator *expr = llvm::cast<clang::BinaryOperator>(stmt);
    switch (expr->getOpcode()) {
    case clang::BinaryOperatorKind::BO_Add:
    case clang::BinaryOperatorKind::BO_Mul:
    case clang::BinaryOperatorKind::BO_Sub:
    case clang::BinaryOperatorKind::BO_AddAssign:
    case clang::BinaryOperatorKind::BO_SubAssign:
    case clang::BinaryOperatorKind::BO_MulAssign: {
      const clang::Expr *lhs = extract_simple_value_expr(expr->getLHS());
      const clang::Expr *rhs = extract_simple_value_expr(expr->getRHS());
      bool lhs_result, rhs_result;
      lhs_result = rhs_result = true;
      if (lhs) {
        OperationDetection od(manager, call_graph, ctx.function);
        od.set_operation<BranchOperation>();
        od.set_use_point(ctx.block, ctx.stmt_it);
        dispatch_set_value_holder(od.use_point(), lhs);
        lhs_result = od.detect(OperationDetection::Forall);
      }
      if (rhs) {
        OperationDetection od(manager, call_graph, ctx.function);
        od.set_operation<BranchOperation>();
        od.set_use_point(ctx.block, ctx.stmt_it);
        dispatch_set_value_holder(od.use_point(), rhs);
        rhs_result = od.detect(OperationDetection::Forall);
      }
      if (!(lhs_result || rhs_result)) {
        DefectType dt = DefectType::IntegerOverflow;
        reportJson(rhs, ctx, dt);
        llvm::errs() << "Defect8 Arithmetic Overflow\n";
        expr->dump();
        llvm::errs() << '\n';
        return true;
      }
      break;
    }
    default:
      break;
    }
  } else if (stmt->getStmtClass() == clang::Stmt::UnaryOperatorClass) {
    const clang::UnaryOperator *expr = llvm::cast<clang::UnaryOperator>(stmt);
    switch (expr->getOpcode()) {
    case clang::UnaryOperatorKind::UO_PostInc:
    case clang::UnaryOperatorKind::UO_PostDec:
    case clang::UnaryOperatorKind::UO_PreInc:
    case clang::UnaryOperatorKind::UO_PreDec: {
      const clang::Expr *sub = extract_simple_value_expr(expr->getSubExpr());
      if (sub) {
        OperationDetection od(manager, call_graph, ctx.function);
        od.set_operation<BranchOperation>();
        od.set_use_point(ctx.block, ctx.stmt_it);
        dispatch_set_value_holder(od.use_point(), sub);
        if (!od.detect(OperationDetection::Forall)) {
          DefectType dt = DefectType::IntegerOverflow;
          reportJson(sub, ctx, dt);
          llvm::errs() << "Defect8 Arithmetic Overflow\n";
          expr->dump();
          llvm::errs() << '\n';
          return true;
        }
      }
      break;
    }
    default:
      break;
    }
  }
  return false;
}
void DetectMe::detect_null_dereference(const clang::UnaryOperator *expr,
                                       WalkContext &ctx) {
  if (expr->getOpcode() == clang::UnaryOperatorKind::UO_Deref) {
    const clang::Expr *target = extract_simple_value_expr(expr->getSubExpr());
    if (target) {
      OperationDetection od(manager, call_graph, ctx.function);
      od.set_operation<BranchOperation>();
      od.set_use_point(ctx.block, ctx.stmt_it);
      dispatch_set_value_holder(od.use_point(), target);
      if (!od.detect(OperationDetection::Forall)) {
        DefectType dt = DefectType::NullPointerDereference;
        reportJson(target, ctx, dt);
        llvm::errs() << "Defect9 Null Dereference\n";
        expr->dump();
        llvm::errs() << '\n';
      }
    }
  }
}
void DetectMe::detect_use_after_free(const clang::UnaryOperator *expr,
                                     WalkContext &ctx) {
  if (expr->getOpcode() == clang::UnaryOperatorKind::UO_Deref) {
    const clang::Expr *target = extract_simple_value_expr(expr->getSubExpr());
    if (target) {
        OperationDetection od(manager, call_graph, ctx.function);
        od.set_operation<CallOperation>("free", 0);
        od.set_use_point(ctx.block, ctx.stmt_it);
        dispatch_set_value_holder(od.use_point(), target);
        if (od.detect(OperationDetection::Exists)) {
          DefectType dt = DefectType::UseAfterFree;
          reportJson(target, ctx, dt);
          llvm::errs() << "Defect10 Use After Free\n";
          expr->dump();
          llvm::errs() << '\n';
      }
    }
  }
}
// void DetectMe::detect_illegal_element_access(); // unknown
// void DetectMe::detect_illegal_file_path(); // unknown